{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "funded-italic",
   "metadata": {},
   "source": [
    "## Treatment effects under conditional independence\n",
    "\n",
    "Author: Josef Perktold\n",
    "\n",
    "This notebook illustrates the basic usage of the new treatment effect functionality in statsmodels.\n",
    "\n",
    "The main class is `statsmodels.treatment.treatment_effects.TreatmentEffect`.\n",
    "\n",
    "\n",
    "This class estimates treatment effect and potential outcome using 5 different methods, ipw, ra, aipw, aipw-wls, ipw-ra. The last three methods require both a treatment or selection model and an outcome model.\n",
    "Standard errors and inference are based on the joint GMM representation of selection or treatment model, outcome model and effect functions. The approach for inference follows Stata, however Stata support a wider range of models.\n",
    "Estimation and inference are valid under conditional independence or ignorability.\n",
    "\n",
    "The outcome model is currently limited to a linear model based on OLS.\n",
    "Treatment is currently restricted to binary treatment which can be either Logit or Probit.\n",
    "\n",
    "The example follows Cattaneo."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adult-capture",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from numpy.testing import assert_allclose\n",
    "from statsmodels.discrete.discrete_model import Probit\n",
    "from statsmodels.regression.linear_model import OLS\n",
    "from statsmodels.treatment.tests.results import results_teffects as res_st\n",
    "from statsmodels.treatment.treatment_effects import TreatmentEffect\n",
    "\n",
    "# Load data for example\n",
    "cur_dir = os.path.abspath(os.path.dirname(res_st.__file__))\n",
    "file_name = \"cataneo2.csv\"\n",
    "file_path = os.path.join(cur_dir, file_name)\n",
    "dta_cat = pd.read_csv(file_path)\n",
    "\n",
    "methods = [\"ra\", \"ipw\", \"aipw\", \"aipw_wls\", \"ipw_ra\"]\n",
    "methods_st = [\n",
    "    (\"ra\", res_st.results_ra),\n",
    "    (\"ipw\", res_st.results_ipw),\n",
    "    (\"aipw\", res_st.results_aipw),\n",
    "    (\"aipw_wls\", res_st.results_aipw_wls),\n",
    "    (\"ipw_ra\", res_st.results_ipwra),\n",
    "]\n",
    "\n",
    "# allow wider display of data frames\n",
    "pd.set_option(\"display.width\", 500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e76526",
   "metadata": {},
   "outputs": [],
   "source": [
    "dta_cat.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "talented-bidding",
   "metadata": {},
   "source": [
    "### Create TreatmentEffect instance and compute ipw\n",
    "\n",
    "The TreatmentEffect class requires \n",
    "- a OLS model instance for the outcome model, \n",
    "- a results instance of the selection model and \n",
    "- a treatment indicator variable.\n",
    "\n",
    "In the following example we use Probit as the selection model. Using Logit is also supported.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "controlling-literacy",
   "metadata": {},
   "outputs": [],
   "source": [
    "# treatment selection model\n",
    "formula = \"mbsmoke_ ~ mmarried_ + mage + mage2 + fbaby_ + medu\"\n",
    "res_probit = Probit.from_formula(formula, dta_cat).fit()\n",
    "\n",
    "# outcome model\n",
    "formula_outcome = \"bweight ~ prenatal1_ + mmarried_ + mage + fbaby_\"\n",
    "mod = OLS.from_formula(formula_outcome, dta_cat)\n",
    "\n",
    "# treatment indicator variable\n",
    "tind = np.asarray(dta_cat[\"mbsmoke_\"])\n",
    "\n",
    "teff = TreatmentEffect(mod, tind, results_select=res_probit)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecbd49a4",
   "metadata": {
    "slideshow": {
     "slide_type": "fragment"
    }
   },
   "source": [
    "After creating the TreatmentEffect instance, we can call any of the 5 methods to compute potential outcomes, POM0, POM1, and average treatment effect, ATE. POM0 is the potential outcome for the no treatment group, POM1 is the potential outcome for the treatment group, treatment effect is POM1 - POM0.\n",
    "\n",
    "For example `teff.ipw()` computes POM and ATE using inverse probability weighting. The probability of treatment is also commonly called the propensity score. The `summary` of the estimation includes standard errors and confidence interval for POM and ATE.\n",
    "\n",
    "\n",
    "Standard errors and other inferential statistics are based on the Generalized Method of Moments (GMM) representation of the selection and outcome models and the moment conditions for the results statistic.\n",
    "Method `ipw` uses the selection model but not the outcome model.\n",
    "Method `ra` uses the outcome model but not the selection model.\n",
    "The doubly robust estimators `aipw`, `aipw-wls`, `ipw-ra` include both selection and outcome models, where at least one of those two has to be correctly specified to get consistent estimates of the treatment effect.\n",
    "The moment conditions for the target variables, POM0, POM1, and ATE are based on POM0 and ATE. The remaining POM1 is computed as a linear combination of POM0 and ATE.\n",
    "\n",
    "The internal gmm results are attached to the treatment results as `results_gmm`.\n",
    "\n",
    "By default the treatment effect methods computes average treatment effect, where average is take over the sample observations.\n",
    "Option `effect_group` can be used to compute either average treatment effect on the treated, ATT, using `effect_group=1` or average treatment effect on the non-treated using `effect_group=0`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "former-mirror",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = teff.ipw()\n",
    "res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "metropolitan-burns",
   "metadata": {},
   "outputs": [],
   "source": [
    "res.summary_frame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "documentary-invasion",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(res.results_gmm.summary())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9410bb5",
   "metadata": {},
   "source": [
    "**average treatment effect on the treated**\n",
    "\n",
    "see more below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d00a2846",
   "metadata": {},
   "outputs": [],
   "source": [
    "teff.ipw(effect_group=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "464fbfec",
   "metadata": {},
   "source": [
    "**average treatment effect on the untreated**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df5a0f1d",
   "metadata": {},
   "outputs": [],
   "source": [
    "teff.ipw(effect_group=0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3071337",
   "metadata": {},
   "source": [
    "Other methods to compute ATE work in the same or similar way as for `ipw` for example regression adjustment `ra` and double robust `ipw_ra`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0795f217",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_ra = teff.ra()\n",
    "res_ra"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "vocational-crawford",
   "metadata": {},
   "outputs": [],
   "source": [
    "res_ra.summary_frame()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "posted-basin",
   "metadata": {},
   "outputs": [],
   "source": [
    "ra2 = teff.ipw_ra(effect_group=1, return_results=True)\n",
    "ra2.summary_frame()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "banned-rogers",
   "metadata": {},
   "source": [
    "## All methods in TreatmentEffect\n",
    "\n",
    "The following computes and prints ATE and POM for all methods.\n",
    "(We include the call to TreatmentEffect as a reminder.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "innovative-garage",
   "metadata": {},
   "outputs": [],
   "source": [
    "teff = TreatmentEffect(mod, tind, results_select=res_probit)\n",
    "\n",
    "for m in methods:\n",
    "    res = getattr(teff, m)()\n",
    "    print(\"\\n\", m)\n",
    "    print(res.summary_frame())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "broadband-export",
   "metadata": {},
   "source": [
    "## Results in Stata\n",
    "\n",
    "The results in statsmodels are very close to the results in Stata because both packages use the same approach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "compact-pixel",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m, st in methods_st:\n",
    "    print(\"\\n\", m)\n",
    "    res = pd.DataFrame(\n",
    "        st.table[:2, :6], index=[\"ATE\", \"POM0\"], columns=st.table_colnames[:6]\n",
    "    )\n",
    "    print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "physical-newman",
   "metadata": {},
   "source": [
    "### Treatment effects without inference\n",
    "\n",
    "It is possible to compute POM and ATE without computing standard errors and inferential statistics. In this case the GMM model is not computed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "preceding-dividend",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in methods:\n",
    "    print(\"\\n\", m)\n",
    "    res = getattr(teff, m)(return_results=False)\n",
    "    print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "personal-marble",
   "metadata": {},
   "source": [
    "## Treatment effect on the treated\n",
    "\n",
    "Treatment effects on subgroups are not available for `aipw` and `aipw-wls`.\n",
    "\n",
    "`effect_group` choses the group for which treatement effect and potential outcomes are computed\n",
    "Options are \n",
    "\"all\" for sample average treatment effect, \n",
    "`1` for average treatment effect on the treated and \n",
    "`0` for average treatment effect on the untreated.\n",
    "\n",
    "Note: The row labels in the pandas dataframe, POM and ATE, are the same even for treatment effect on subgroups."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "lyric-producer",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in methods:\n",
    "    if m.startswith(\"aipw\"):\n",
    "        continue\n",
    "    res = getattr(teff, m)(effect_group=1)\n",
    "    print(\"\\n\", m)\n",
    "    print(res.summary_frame())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "speaking-functionality",
   "metadata": {},
   "source": [
    "### Treatment effect on the untreated\n",
    "\n",
    "Similar to ATT, we can compute average treatment effect on the untreated by using option `effect_group=0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "southern-startup",
   "metadata": {},
   "outputs": [],
   "source": [
    "for m in methods:\n",
    "    if m.startswith(\"aipw\"):\n",
    "        # not available\n",
    "        continue\n",
    "    res = getattr(teff, m)(effect_group=0)\n",
    "    print(\"\\n\", m)\n",
    "    print(res.summary_frame())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a88cd1ff",
   "metadata": {},
   "source": [
    "The docstring for the TreatmentEffect class and it's methods can be obtained using help\n",
    "\n",
    "`help(teff)`"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
